package com.study.talk.ai

import com.study.talk.BuildConfig
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage
import com.volcengine.ark.runtime.service.ArkService
import io.reactivex.disposables.Disposable
import io.reactivex.schedulers.Schedulers
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import okhttp3.ConnectionPool
import okhttp3.Dispatcher
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean

object AiModel {

    private const val API_KEY = BuildConfig.AI_API_KEY

    private const val CHAT_MODEL = BuildConfig.CHAT_MODEL
    private const val VISION_MODEL = BuildConfig.VISION_MODEL

    private const val TAG = "AiModel"

    private val service =
        ArkService.builder().baseUrl("https://ark.cn-beijing.volces.com/api/v3").apiKey(API_KEY)
            .callbackExecutor(Executors.newSingleThreadExecutor())
            .connectionPool(ConnectionPool(5, 1, TimeUnit.SECONDS)).dispatcher(Dispatcher()).build()

    fun sendMsg(
        msg: List<ChatMessage>,
        start: () -> Unit,
        subscribe: (response: String) -> Unit,
        finish: () -> Unit,
        error: (Throwable) -> Unit = {}
    ): Disposable {
        start()
        val buffer = StringBuilder()
        val doing = AtomicBoolean(false)
        val request = ChatCompletionRequest.builder().model(CHAT_MODEL).messages(msg).build()
        return service.streamChatCompletion(request).doOnError(error).subscribeOn(Schedulers.io())
            .observeOn(Schedulers.io()).doOnComplete(finish).subscribe {
                if (it.choices.isNotEmpty()) {
                    buffer.append(it.choices[0].message.content.toString())
                    if (!doing.getAndSet(true)) {
                        CoroutineScope(Dispatchers.IO).launch {
                            delay(500)
                            withContext(Dispatchers.Main) {
                                subscribe(buffer.toString())
                            }
                            doing.set(false)
                        }
                    }
                }
            }
    }

    fun analyzeImage(
        msg: List<ChatMessage>,
        start: () -> Unit,
        subscribe: (response: String) -> Unit,
        finish: () -> Unit,
        error: (Throwable) -> Unit = {}
    ): Disposable {
        start()
        val buffer = StringBuilder()
        val doing = AtomicBoolean(false)
        val request = ChatCompletionRequest.builder().model(VISION_MODEL).messages(msg).build()
        return service.streamChatCompletion(request).doOnError(error).subscribeOn(Schedulers.io())
            .observeOn(Schedulers.io()).doOnComplete(finish).subscribe {
                if (it.choices.isNotEmpty()) {
                    buffer.append(it.choices[0].message.content.toString())
                    if (!doing.getAndSet(true)) {
                        CoroutineScope(Dispatchers.IO).launch {
                            delay(500)
                            withContext(Dispatchers.Main) {
                                subscribe(buffer.toString())
                            }
                            doing.set(false)
                        }
                    }
                }
            }
    }
}